using numba to speed up mean-squared displacement calculations
Writing a faster mean-squared-displacement function¶
I'm trying to calculate the Mean Squared Displacement (MSD) of a particle trajectory. In reality, I have an array of positions for numtrj particles and numpts timepoints and dim dimensions: pos[numtrj, numpts, dim]. I think my question has the same answer if I just have the x trajectory of a single particle, though.
In case you haven't done MSD calculations before, there's one cute way in which you get extra information out of a trajectory. Say I have just five time points, and my positions are
In [83]: x
Out[83]: array([ 0. , 1.74528704, 1.59639865, 2.59976219, 3.70852457])
You could just get squared displacement by looking at x**2. However, you could also say that you have 4 different values for the displacement at one timestep:
x[1] - x[0], x[2] - x[1], x[3] - x[2], x[4] - x[3]
Similarly, three values for displacement at two timesteps: x[2:] - x[:-2]. Etc. So, the way I'm calculating MSD at the moment is:
def msd_1d(x):
result = np.zeros_like(x)
for i in range(1,len(x)):
result[i] = np.average((x[i:] - x[:-i])**2)
return result
or
def get_msd_traj(pos):
result = np.zeros_like(pos)
for i in range(1,pos.shape[1]):
result[:,i,:] = np.average((pos[:,i:,:] - pos[:,:-i,:])**2,axis=1)
return result
(side note: often the data comes indexed like pos[numpts, numtrj, dim] for molecular dynamics trajectories, but that doesn't change anything here)
So, I asked Joshua Adelman if he had any quick thoughts.

Your IPyNB class could be made a lot more readable by using a custom JSON encoder, e.g. extending JSONEncoder http://docs.python.org/2/library/json.html#json.JSONEncoder. That way, you could store the preamble/afterward as a Python dict and just call json.dumps(..., default=MyJSONEncoder) inside tostring().
Updated above, after our email conversation.